from captcha.image import ImageCaptcha
import matplotlib.pyplot as plt
import random
import string
import os
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor, as_completed


def generate(show=False, n_len=1):    
    characters = string.digits + string.ascii_uppercase
    characters = ''.join([cha for cha in characters if cha not in 'O01'])
    width, height, n_class = 160, 60, len(characters)
    generator = ImageCaptcha(width=width, height=height)

    random_str = ''.join([random.choice(characters) for j in range(n_len)])
    img = generator.generate_image(random_str)
    if show:
        plt.imshow(img)
        plt.title(random_str)
        plt.show()
    return random_str, img


def task(params):
    idx, n_len, out_dir = params
    if not n_len:
        n_len = random.randint(1,7)
    label, img = generate(n_len=n_len)
    label_dir = out_dir  # os.path.join(out_dir, label)
    if not os.path.isdir(label_dir):
        os.makedirs(label_dir)
    path = os.path.join(label_dir, f'{label}.{idx}.png')
    img.save(path)
    return n_len


def batch_generate(out_dir='captcha3/train', batch_size=1000, n_len=None):
    with ProcessPoolExecutor() as pool:
        tasks = [pool.submit(task, [i, n_len, out_dir]) for i in range(batch_size)]
        print('任务提交完毕，等待执行结果...')
        for _ in tqdm(as_completed(tasks), total=batch_size):
            pass


if __name__ == '__main__':
    # generate(show=True)
    batch_generate(out_dir=r'C:/Users/yczha/Documents/torch_test/dataset/captcha/train', batch_size=10000, n_len=None)
    batch_generate(out_dir=r'C:/Users/yczha/Documents/torch_test/dataset/captcha/test', batch_size=1000, n_len=None)
    # batch_generate(out_dir='captcha4/test', batch_size=5000, n_len=4)
    # batch_generate(out_dir='captcha4/valid', batch_size=5000, n_len=4)

